# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/9/24 17:06
# @Author  : Dell
# @File    : split_question.py
# @Software: PyCharm
# @Desc    :问题分解：把复杂/宏观的问题拆解成多个小问题，然后自动进行多轮对话，最后总结回答最初的问题
import json
from typing import Union
from operator import itemgetter
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.documents import Document

llm = ChatOpenAI(model_name='qwen-plus', base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", temperature=0.5,
                 api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37")
# --------------文档基础处理，索引和检索开始-----------------------
# 1、加载文本
loader = TextLoader(file_path="xiyouji.txt",encoding="utf-8")
blog_docs = loader.load()
# 2、分割文档
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=300,chunk_overlap=50)
splits = text_splitter.split_documents(blog_docs)
# 3、构建一个支持语义搜索的向量数据库（数据库在内存中）
embd=DashScopeEmbeddings(
    model="text-embedding-v1", dashscope_api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37"
)
vectorstore = Chroma.from_documents(documents=splits,embedding=embd)
#4、检索
# 将一个向量数据库（vectorstore）转换为一个检索器（retriever），并设置检索参数：返回最相似的5个结果。
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
# --------------文档基础处理，索引和检索结束-----------------------
# ------------------分解问题开始--------------------------------
# 分解提示词
template = """您是一位乐于助人的助手，可以生成与输入问题相关的多个子问题。\n
目标是将输入分解为一组可以单独回答的子问题/子问题。\n
生成与以下问题相关的多个搜索查询：{question} \n
输出（5 个查询）："""
prompt_decomposition = ChatPromptTemplate.from_template(template)
# 构造一个分解链
generate_queries_decomposition = (prompt_decomposition | llm | StrOutputParser() | (lambda x: x.split("\n")))

question = "根据文章，孙悟空是怎么一步步发展到今天的？"
questions = generate_queries_decomposition.invoke({"question":question})
print(f"拆解后的问题清单：{questions}")
# ------------------分解问题结束--------------------------------

# ------------------对问题进行整合开始--------------------------------

template = """这是您需要回答的问题：

\n --- \n {question} \n --- \n

以下是任何可用的背景问题 + 答案组合：

\n --- \n {q_a_pairs} \n --- \n

以下是与问题相关的其他上下文：

\n --- \n {context} \n --- \n

使用上述上下文和任何背景问题 + 答案组合来回答问题：\n {question}
"""

decomposition_prompt = ChatPromptTemplate.from_template(template)
# 把所有东西都串联起来
def format_qa_pair(question, answer):
    """格式化 Q 和 A 对"""

    formatted_string = ""
    formatted_string += f"Question: {question}\nAnswer: {answer}\n\n"
    return formatted_string.strip()


q_a_pairs = ""
# 循环进行Q&A
for q in questions:
    # 构建rag链 itemgetter()函数是从字典或对象中提取特定字段的高效工具
    rag_chain = (
    {"context": itemgetter("question") | retriever,"question": itemgetter("question"),"q_a_pairs": itemgetter("q_a_pairs")}
    | decomposition_prompt
    | llm
    | StrOutputParser())

    answer = rag_chain.invoke({"question":q,"q_a_pairs":q_a_pairs})
    q_a_pair = format_qa_pair(q,answer)
    q_a_pairs = q_a_pairs + "\n---\n"+  q_a_pair

    print(f"打印{q}的答案：{answer}\n")

# ------------------对问题进行整合开始--------------------------------